package net.aiion.weave.impl.bytecode;

import java.lang.annotation.Annotation;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;

import net.aiion.weave.spi.bytecode.BytecodeException;
import net.aiion.weave.spi.bytecode.BytecodeStore;
import net.aiion.weave.spi.bytecode.BytecodeType;

public class BytecodeStoreImpl implements BytecodeStore
{
    private Map<String, BytecodeType> typesByName;
    // private Map<String, Collection<BytecodeType>> typesBySuperName;
    private Map<String, Collection<BytecodeType>> typesByInterfaceName;
    private Map<String, Collection<BytecodeType>> typesByAnnotationName;

    private Map<String, byte[]> bytecodeClasses;
    private BytecodeLoader bytecodeLoader;

    public BytecodeStoreImpl() {
        typesByName = new HashMap<>();
        // typesBySuperName = new HashMap<>();
        typesByInterfaceName = new HashMap<>();
        typesByAnnotationName = new HashMap<>();

        bytecodeClasses = new HashMap<>();
        bytecodeLoader = new BytecodeLoader(this, Thread.currentThread().getContextClassLoader());
    }

    @Override
    public void registerType(final BytecodeType aBytecodeType) {
        typesByName.put(aBytecodeType.getTypeName(), aBytecodeType);

        // if (!typesBySuperName.containsKey(aBytecodeType.getSuperTypeName())) {
        // typesBySuperName.put(aBytecodeType.getSuperTypeName(), new ArrayList<BytecodeType>());
        // }
        //
        // typesBySuperName.get(aBytecodeType.getSuperTypeName()).add(aBytecodeType);

        for (String interfaceName : aBytecodeType.getInterfaceNames()) {
            if (!typesByInterfaceName.containsKey(interfaceName)) {
                typesByInterfaceName.put(interfaceName, new ArrayList<BytecodeType>());
            }

            typesByInterfaceName.get(interfaceName).add(aBytecodeType);
        }

        for (String annotationName : aBytecodeType.getAnnotationNames()) {
            if (!typesByAnnotationName.containsKey(annotationName)) {
                typesByAnnotationName.put(annotationName, new ArrayList<BytecodeType>());
            }

            typesByAnnotationName.get(annotationName).add(aBytecodeType);
        }
    }

    @Override
    public void registerClass(final String aName, final byte[] aBytes) {
        bytecodeClasses.put(aName.replace('/', '.'), aBytes);
    }

    public byte[] removeClass(final String aName) {
        return bytecodeClasses.remove(aName.replace('/', '.'));
    }

    @Override
    public BytecodeType findTypeByName(final String aName) {
        return typesByName.get(aName);
    }

    @Override
    public Class<?> findClassByName(final String aName) {
        try {
            return bytecodeLoader.loadClass(aName);
        } catch (final ClassNotFoundException e) {
            return null;
        }
    }

    @Override
    public Collection<BytecodeType> findTypesByInterfaceName(final String aInterfaceName) {
        if (!typesByInterfaceName.containsKey(aInterfaceName)) {
            return Collections.emptyList();
        }

        final Collection<BytecodeType> subTypes = new ArrayList<>();
        for (BytecodeType subType : typesByInterfaceName.get(aInterfaceName)) {
            subTypes.add(subType);
            // XXX subTypes.addAll(findTypesBySuperType(subType));
        }
        return subTypes;
    }

    @Override
    public Collection<BytecodeType> findTypesByInterfaceClass(final Class<?> aInterfaceClass) {
        return findTypesByInterfaceName(aInterfaceClass.getName());
    }

    @Override
    public Collection<Class<?>> findClassesByInterfaceName(final String aInterfaceName) {
        return loadClasses(findTypesByInterfaceName(aInterfaceName));
    }

    @Override
    public <T> Collection<Class<? extends T>> findClassesByInterfaceClass(final Class<T> aInterfaceClass) {
        final Collection<Class<?>> uncheckedClasses = findClassesByInterfaceName(aInterfaceClass.getName());
        final Collection<Class<? extends T>> checkedClasses = new ArrayList<>(uncheckedClasses.size());
        for (Class<?> c : uncheckedClasses) {
            checkedClasses.add(c.asSubclass(aInterfaceClass));
        }
        return checkedClasses;
    }

    @Override
    public Collection<BytecodeType> findTypesByAnnotationName(final String aAnnotationName) {
        if (!typesByAnnotationName.containsKey(aAnnotationName)) {
            return Collections.emptyList();
        }

        return typesByAnnotationName.get(aAnnotationName);
    }

    @Override
    public Collection<BytecodeType> findTypesByAnnotationClass(final Class<? extends Annotation> aAnnotationClass) {
        return findTypesByAnnotationName(aAnnotationClass.getName());
    }

    @Override
    public Collection<Class<?>> findClassesByAnnotationName(final String aAnnotationName) {
        return loadClasses(findTypesByAnnotationName(aAnnotationName));
    }

    @Override
    public Collection<Class<?>> findClassesByAnnotationClass(Class<? extends Annotation> aAnnotationClass) {
        return findClassesByAnnotationName(aAnnotationClass.getName());
    }

    private Collection<Class<?>> loadClasses(final Collection<BytecodeType> aBytecodeTypes) {
        final Collection<Class<?>> classes = new ArrayList<>(aBytecodeTypes.size());
        for (BytecodeType bytecodeType : aBytecodeTypes) {
            classes.add(loadClass(bytecodeType.getTypeName()));
        }
        return classes;
    }

    private Class<?> loadClass(final String aName) {
        try {
            return bytecodeLoader.loadClass(aName);
        } catch (final ClassNotFoundException e) {
            throw new BytecodeException(e);
        }
    }
}
